use crate::ast::*;
use crate::error::{CompilerError, Result};
use super::super::super::TypeEnvironment;
use super::super::super::TypeInference;

impl TypeInference {




    /// Infer types for Result method calls
    pub(in crate::type_check::type_inference) fn infer_result_method_call(
        &mut self,
        result_type: &crate::ast::types::ResultType,
        method_name: &str,
        arguments: &[Expression],
        location: &crate::error::Location,
        env: &mut TypeEnvironment,
    ) -> Result<Type> {
        match method_name {
            "is_ok" => {
                if !arguments.is_empty() {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.is_ok() expects 0 arguments, got {}",
                            arguments.len()
                        ),
                    ));
                }
                Ok(Type::Basic(crate::ast::types::BasicType::Bool))
            }
            "is_err" => {
                if !arguments.is_empty() {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.is_err() expects 0 arguments, got {}",
                            arguments.len()
                        ),
                    ));
                }
                Ok(Type::Basic(crate::ast::types::BasicType::Bool))
            }
            "unwrap" => {
                if !arguments.is_empty() {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.unwrap() expects 0 arguments, got {}",
                            arguments.len()
                        ),
                    ));
                }
                Ok(*result_type.ok_type.clone())
            }
            "expect" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.expect() expects 1 argument (message), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Check that the argument is a string
                let msg_type = self.infer_expression(&arguments[0], env)?;
                self.add_constraint(
                    &msg_type,
                    &Type::Basic(crate::ast::types::BasicType::String),
                    *location,
                )?;
                Ok(*result_type.ok_type.clone())
            }
            "expect_err" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.expect_err() expects 1 argument (message), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Check that the argument is a string
                let msg_type = self.infer_expression(&arguments[0], env)?;
                self.add_constraint(
                    &msg_type,
                    &Type::Basic(crate::ast::types::BasicType::String),
                    *location,
                )?;
                Ok(*result_type.err_type.clone())
            }
            "unwrap_or" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.unwrap_or() expects 1 argument (default value), got {}",
                            arguments.len()
                        ),
                    ));
                }
                let default_type = self.infer_expression(&arguments[0], env)?;
                self.add_constraint(&default_type, &*result_type.ok_type, *location)?;
                Ok(*result_type.ok_type.clone())
            }
            "unwrap_or_else" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.unwrap_or_else() expects 1 argument (function), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer the function type
                let func_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Function(FunctionType {
                    parameter_types,
                    return_type,
                    ..
                }) = func_type
                {
                    if parameter_types.len() != 1 {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Result.unwrap_or_else() function must take exactly 1 argument",
                        ));
                    }
                    // Check that the function's parameter type matches Err type
                    self.add_constraint(&parameter_types[0], &*result_type.err_type, *location)?;

                    // Check that the function's return type matches Ok type
                    let func_return_type = return_type
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                    self.add_constraint(&func_return_type, &*result_type.ok_type, *location)?;

                    Ok(*result_type.ok_type.clone())
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Result.unwrap_or_else() expects a function argument",
                    ))
                }
            }
            "unwrap_err" => {
                if !arguments.is_empty() {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.unwrap_err() expects 0 arguments, got {}",
                            arguments.len()
                        ),
                    ));
                }
                Ok(*result_type.err_type.clone())
            }
            "ok" => {
                if !arguments.is_empty() {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!("Result.ok() expects 0 arguments, got {}", arguments.len()),
                    ));
                }
                Ok(Type::Optional(crate::ast::types::OptionalType {
                    inner_type: result_type.ok_type.clone(),
                    location: *location,
                }))
            }
            "err" => {
                if !arguments.is_empty() {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!("Result.err() expects 0 arguments, got {}", arguments.len()),
                    ));
                }
                Ok(Type::Optional(crate::ast::types::OptionalType {
                    inner_type: result_type.err_type.clone(),
                    location: *location,
                }))
            }
            "map" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.map() expects 1 argument (function), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer the function type
                let func_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Function(FunctionType {
                    parameter_types,
                    return_type,
                    ..
                }) = func_type
                {
                    if parameter_types.len() != 1 {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Result.map() function must take exactly 1 argument",
                        ));
                    }
                    // Check that the function's parameter type matches Ok type
                    self.add_constraint(&parameter_types[0], &*result_type.ok_type, *location)?;

                    // Return type is Result<U, E> where U is the function's return type
                    let new_ok_type = return_type
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                    Ok(Type::Result(crate::ast::types::ResultType {
                        ok_type: new_ok_type,
                        err_type: result_type.err_type.clone(),
                        location: *location,
                    }))
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Result.map() expects a function argument",
                    ))
                }
            }
            "map_err" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.map_err() expects 1 argument (function), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer the function type
                let func_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Function(FunctionType {
                    parameter_types,
                    return_type,
                    ..
                }) = func_type
                {
                    if parameter_types.len() != 1 {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Result.map_err() function must take exactly 1 argument",
                        ));
                    }
                    // Check that the function's parameter type matches Err type
                    self.add_constraint(&parameter_types[0], &*result_type.err_type, *location)?;

                    // Return type is Result<T, F> where F is the function's return type
                    let new_err_type = return_type
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::String)));
                    Ok(Type::Result(crate::ast::types::ResultType {
                        ok_type: result_type.ok_type.clone(),
                        err_type: new_err_type,
                        location: *location,
                    }))
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Result.map_err() expects a function argument",
                    ))
                }
            }
            "map_or" => {
                if arguments.len() != 2 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.map_or() expects 2 arguments (default, function), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer default type
                let default_type = self.infer_expression(&arguments[0], env)?;

                // Infer the function type
                let func_type = self.infer_expression(&arguments[1], env)?;
                if let Type::Function(FunctionType {
                    parameter_types,
                    return_type,
                    ..
                }) = func_type
                {
                    if parameter_types.len() != 1 {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Result.map_or() function must take exactly 1 argument",
                        ));
                    }
                    // Check that the function's parameter type matches Ok type
                    self.add_constraint(&parameter_types[0], &*result_type.ok_type, *location)?;

                    // Check that the function's return type matches default type
                    let func_return_type = return_type
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                    self.add_constraint(&func_return_type, &default_type, *location)?;

                    // Return type is the default type (or function return type, they should match)
                    Ok(default_type)
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Result.map_or() expects a function as second argument",
                    ))
                }
            }
            "map_or_else" => {
                if arguments.len() != 2 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!("Result.map_or_else() expects 2 arguments (default_fn, function), got {}", arguments.len()),
                    ));
                }
                // Infer default_fn type
                let default_fn_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Function(FunctionType {
                    parameter_types: default_params,
                    return_type: default_return,
                    ..
                }) = default_fn_type
                {
                    if default_params.len() != 1 {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Result.map_or_else() default_fn must take exactly 1 argument (error value)",
                        ));
                    }
                    // Check that default_fn's parameter type matches Err type
                    self.add_constraint(&default_params[0], &*result_type.err_type, *location)?;

                    let default_type = default_return
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));

                    // Infer the function type
                    let func_type = self.infer_expression(&arguments[1], env)?;
                    if let Type::Function(FunctionType {
                        parameter_types,
                        return_type,
                        ..
                    }) = func_type
                    {
                        if parameter_types.len() != 1 {
                            return Err(CompilerError::type_error(
                                location.line,
                                location.column,
                                "Result.map_or_else() function must take exactly 1 argument",
                            ));
                        }
                        // Check that the function's parameter type matches Ok type
                        self.add_constraint(&parameter_types[0], &*result_type.ok_type, *location)?;

                        // Check that the function's return type matches default type
                        let func_return_type = return_type
                            .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                        self.add_constraint(&func_return_type, &*default_type, *location)?;

                        // Return type is the default type (or function return type, they should match)
                        Ok(*default_type)
                    } else {
                        Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Result.map_or_else() expects a function as second argument",
                        ))
                    }
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Result.map_or_else() expects a function as first argument",
                    ))
                }
            }
            "and_then" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.and_then() expects 1 argument (function), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer the function type
                let func_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Function(FunctionType {
                    parameter_types,
                    return_type,
                    ..
                }) = func_type
                {
                    if parameter_types.len() != 1 {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Result.and_then() function must take exactly 1 argument",
                        ));
                    }
                    // Check that the function's parameter type matches Ok type
                    self.add_constraint(&parameter_types[0], &*result_type.ok_type, *location)?;

                    // Check that the function returns a Result type
                    let func_return_type = return_type
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                    if let Type::Result(new_result_type) = *func_return_type {
                        // The return type is the function's return Result type
                        Ok(Type::Result(new_result_type))
                    } else {
                        Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Result.and_then() function must return a Result type",
                        ))
                    }
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Result.and_then() expects a function argument",
                    ))
                }
            }
            "or_else" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.or_else() expects 1 argument (function), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer the function type
                let func_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Function(FunctionType {
                    parameter_types,
                    return_type,
                    ..
                }) = func_type
                {
                    if parameter_types.len() != 1 {
                        return Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Result.or_else() function must take exactly 1 argument",
                        ));
                    }
                    // Check that the function's parameter type matches Err type
                    self.add_constraint(&parameter_types[0], &*result_type.err_type, *location)?;

                    // Check that the function returns a Result type
                    let func_return_type = return_type
                        .unwrap_or(Box::new(Type::Basic(crate::ast::types::BasicType::Void)));
                    if let Type::Result(new_result_type) = *func_return_type {
                        // The return type is the function's return Result type
                        Ok(Type::Result(new_result_type))
                    } else {
                        Err(CompilerError::type_error(
                            location.line,
                            location.column,
                            "Result.or_else() function must return a Result type",
                        ))
                    }
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Result.or_else() expects a function argument",
                    ))
                }
            }
            "and" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.and() expects 1 argument (other Result), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer the other Result type
                let other_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Result(other_result_type) = other_type {
                    // Check that error types match
                    self.add_constraint(
                        &other_result_type.err_type,
                        &*result_type.err_type,
                        *location,
                    )?;

                    // Return type is Result<U, E> where U is other's Ok type, E is current's Err type
                    Ok(Type::Result(crate::ast::types::ResultType {
                        ok_type: other_result_type.ok_type,
                        err_type: result_type.err_type.clone(),
                        location: *location,
                    }))
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Result.and() expects a Result argument",
                    ))
                }
            }
            "or" => {
                if arguments.len() != 1 {
                    return Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        format!(
                            "Result.or() expects 1 argument (other Result), got {}",
                            arguments.len()
                        ),
                    ));
                }
                // Infer the other Result type
                let other_type = self.infer_expression(&arguments[0], env)?;
                if let Type::Result(other_result_type) = other_type {
                    // Check that Ok types match
                    self.add_constraint(
                        &other_result_type.ok_type,
                        &*result_type.ok_type,
                        *location,
                    )?;

                    // Return type is Result<T, F> where T is current's Ok type, F is other's Err type
                    Ok(Type::Result(crate::ast::types::ResultType {
                        ok_type: result_type.ok_type.clone(),
                        err_type: other_result_type.err_type,
                        location: *location,
                    }))
                } else {
                    Err(CompilerError::type_error(
                        location.line,
                        location.column,
                        "Result.or() expects a Result argument",
                    ))
                }
            }
            _ => Err(CompilerError::type_error(
                location.line,
                location.column,
                format!("Unknown Result method: {}", method_name),
            )),
        }
    }
}
